package org.wheel;

import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.ExpressionVisitorAdapter;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;

import java.util.ArrayList;
import java.util.List;
import java.util.Set;

public class SQLPermissionChecker {
    private final Set<String> allowedTables;
    private final Set<String> allowedColumns;
    private boolean hasPermission = true;
    private List<String> violations = new ArrayList<>();

    public SQLPermissionChecker(Set<String> allowedTables, Set<String> allowedColumns) {
        this.allowedTables = allowedTables;
        this.allowedColumns = allowedColumns;
    }

    public boolean checkPermission(String sql) throws JSQLParserException {
        Statement statement = CCJSqlParserUtil.parse(sql);
        
        if (statement instanceof Select) {
            Select select = (Select) statement;
            if (select.getSelectBody() instanceof PlainSelect) {
                PlainSelect plainSelect = (PlainSelect) select.getSelectBody();
                checkPlainSelect(plainSelect);
            }
        } else {
            violations.add("只支持SELECT语句的权限检查");
            hasPermission = false;
        }
        
        return hasPermission;
    }

    private void checkPlainSelect(PlainSelect plainSelect) {
        // 检查FROM子句中的表
        if (plainSelect.getFromItem() instanceof Table) {
            Table table = (Table) plainSelect.getFromItem();
            checkTablePermission(table);
        }

        // 检查WHERE子句中的列
        if (plainSelect.getWhere() != null) {
            ExpressionVisitorAdapter<Void> expressionVisitor = new ExpressionVisitorAdapter<Void>() {
                @Override
                public <S> Void visit(Column column, S context) {
                    checkColumnPermission(column);
                    return null;
                }
            };
            plainSelect.getWhere().accept(expressionVisitor, null);
        }

        // 检查SELECT列
        if (plainSelect.getSelectItems() != null) {
            for (var item : plainSelect.getSelectItems()) {
                // 直接检查表达式中的列
                // 假设item有getExpression方法
                try {
                    java.lang.reflect.Method getExpression = item.getClass().getMethod("getExpression");
                    Object expression = getExpression.invoke(item);
                    if (expression instanceof net.sf.jsqlparser.expression.Expression) {
                        // 对于所有类型的表达式，使用带泛型的表达式访问者统一处理
                        ExpressionVisitorAdapter<Void> expressionVisitor = new ExpressionVisitorAdapter<Void>() {
                            @Override
                            public <S> Void visit(Column column, S context) {
                                System.out.println("ExpressionVisitorAdapter访问列: " + column.getColumnName());
                                checkColumnPermission(column);
                                return null;
                            }
                        };
                        ((net.sf.jsqlparser.expression.Expression) expression).accept(expressionVisitor, null);
                    }
                } catch (Exception e) {
                    // 如果无法获取表达式，跳过此项
                }
            }
        }
    }

    private void checkTablePermission(Table table) {
        String tableName = table.getName().toLowerCase();
        if (!allowedTables.contains(tableName)) {
            violations.add("无权限访问表: " + table.getName());
            hasPermission = false;
        }
    }

    private void checkColumnPermission(Column column) {
        String columnName = column.getColumnName().toLowerCase();
        if (!allowedColumns.contains(columnName)) {
            violations.add("无权限访问列: " + columnName);
            hasPermission = false;
        }
    }

    public boolean hasPermission() {
        return hasPermission;
    }

    public List<String> getViolations() {
        return violations;
    }
}
